package org.example

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.expressions.{Literal, Multiply}
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.rules.Rule

class MultiplyOptimizationRule(spark: SparkSession) extends Rule[LogicalPlan] {
  logWarning("MultiplyOptimizationRule start")

  def apply(plan: LogicalPlan): LogicalPlan = plan transformAllExpressions {
    case Multiply(left, right, failOnError) if right.isInstanceOf[Literal]
      && right.asInstanceOf[Literal]
      .value.asInstanceOf[Int] == 1 => left
    case Multiply(left, right, failOnError) if left.isInstanceOf[Literal] &&
      left.asInstanceOf[Literal].value.asInstanceOf == 1 => right
  }
}